package net.maku.jw.util;

import cn.hutool.core.io.resource.ResourceUtil;

import java.awt.*;
import java.awt.image.BufferedImage;
import java.util.ArrayList;
import java.util.List;

/**
 * OCR识别工具类
 * @author zhangzhiqin
 * @since 2022/5/9 21:34
 */
public class OcrUtils {

    private OcrUtils() {}

    public static char[][] LIBRARY;

    /**
     * 初始化操作
     */
    public static void init() {
        System.out.println("正在初始化验证码智能识别训练模型...");
        String content = ResourceUtil.readUtf8Str("library.model");
        String[] split = content.split("\\n");
        LIBRARY = new char[split.length][split[0].length()];

        int temp = 0;
        for (String item : split) {
            LIBRARY[temp++] = item.toCharArray();
        }
        System.out.println("验证码识别模型初始化成功");
    }

    /**
     * 获取验证码字符
     * @param image 图片信息
     * @return 验证码
     */
    public static String discernImg(BufferedImage image) {
        StringBuilder sb = new StringBuilder();
        List<BufferedImage> imgList = splitImg(image);

        for (BufferedImage bufferedImage : imgList) {
            // 二值化
            char[] binary = binary(bufferedImage);
            // 识别
            sb.append(discern(binary));
        }
        return sb.toString();
    }

    /**
     * 识别单个字符
     * @param chs 图片二值信息
     * @return 字符
     */
    public static String discern(char[] chs) {
        int target = 0;
        double temp = 0;

        for (int i = 0; i < LIBRARY.length; ++i) {
            double libCount = 1;
            int chCount = 0;

            for (int j = 0; j < LIBRARY[i].length - 1; ++j) {
                if (LIBRARY[i][j] == '1') {
                    ++libCount;
                    if (LIBRARY[i][j] == chs[j]) {
                        ++chCount;
                    }
                }
            }

            if (chCount / libCount > temp) {
                temp = chCount / libCount;
                target = i;
            }

        }

        return String.valueOf(LIBRARY[target][LIBRARY[target].length - 1]);
    }

    /**
     * 切割并黑白化
     * @param bufferedImage 图片信息
     * @return 处理后的图片
     */
    public static List<BufferedImage> splitImg(BufferedImage bufferedImage) {
        List<BufferedImage> imgList = new ArrayList<>();

        imgList.add(removeBackground(bufferedImage.getSubimage(4, 10, 17, 24)));
        imgList.add(removeBackground(bufferedImage.getSubimage(22, 10, 17, 24)));
        imgList.add(removeBackground(bufferedImage.getSubimage(41, 10, 17, 24)));
        imgList.add(removeBackground(bufferedImage.getSubimage(59, 10, 17, 24)));

        return imgList;
    }

    /**
     * 二值化处理
     * @param image 图片信息
     * @return 二值数组
     */
    public static char[] binary(BufferedImage image) {
        int width = image.getWidth();
        int height = image.getHeight();

        char[] chars = new char[width * height];

        int temp = 0;
        for (int i = 0; i < height; ++i) {
            for (int j = 0; j < width; ++j) {
                chars[temp++] = isBackground(image.getRGB(j, i)) ? '0' : '1';
            }
        }
        return chars;
    }

    /**
     * 移除背景并黑白处理
     * @param image 图片信息
     * @return 图片信息
     */
    public static BufferedImage removeBackground(BufferedImage image) {
        for (int x = 0; x < image.getWidth(); ++x) {
            for (int y = 0; y < image.getHeight(); ++y) {
                // 判断是否为字体部分
                if (isBackground(image.getRGB(x, y))) {
                    image.setRGB(x, y, Color.WHITE.getRGB());
                } else {
                    image.setRGB(x, y, Color.BLACK.getRGB());
                }
            }
        }
        return image;
    }

    /**
     * 判断是否背景色
     * @param rgb rgb
     * @return true/false
     */
    public static boolean isBackground(int rgb) {
        Color color = new Color(rgb);
        // 获取三原色平均值
        int avg = (color.getRed() + color.getGreen() + color.getBlue()) / 3;
        return avg > 110;
    }

}
